public class RBTree {
    static class RBTreeNode{
        public RBTreeNode left;
        public RBTreeNode right;
        public RBTreeNode parent;
        public int val;
        public COLOR color;

        public RBTreeNode(int val){
            this.val = val;
            // 节点默认的颜色为红色，因为如果插入黑色节点，会导致不同路径上黑色节点的数量不一样多，
            // 这样还需要插入额外的黑色节点，才能保持红黑树的性质
            this.color = COLOR.RED;
        }
    }

    public RBTreeNode root;

    public boolean  insert(int val){
        RBTreeNode node = new RBTreeNode(val);

        if(this.root == null){
            this.root = node;
            this.root.color = COLOR.BLACK;
            return true;
        }

        RBTreeNode cur = root;
        RBTreeNode parent = null;

        while(cur != null){
            if(cur.val < val){
                parent = cur;
                cur = cur.right;
            }else if(cur.val > val){
                parent = cur;
                cur = cur.left;
            }else{
                return false;
            }
        }

        // cur == null，找到了应该插入的位置
        if(parent.val < val){
            parent.right = node;
        }else{
            parent.left = node;
        }

        node.parent = parent;
        cur = node;

        // 修改颜色
        while(parent != null && parent.color == COLOR.RED){
            RBTreeNode grandFather = parent.parent;

            if(grandFather.left == parent){
                RBTreeNode uncle = grandFather.right;

                // 情况 1
                if(uncle != null && uncle.color == COLOR.RED){
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;

                    cur = grandFather;
                    parent = cur.parent;
                }else{
                    // parent.color == COLOR.BLANK || uncle.color == COLOR.BLANK
                    // 情况 3：
                    if(parent.right == cur){
                        // 左旋
                        rotateLeft(parent);
                        RBTreeNode tmp = parent;
                        parent = cur;
                        cur = tmp;
                    }

                    // 情况 2：parent.color == COLOR.RED grandFather.color = COLOR.BLANK uncle.color == COLOR.BLANK
                    rotateRight(grandFather);
                    parent.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                }
            }else{
                // grandFather.right = parent
                RBTreeNode uncle = grandFather.left;

                // 情况 1：
                if(uncle != null && uncle.color == COLOR.RED){
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;

                    cur = grandFather;
                    parent = cur.parent;
                }else{
                    // parent.color == COLOR.BLANK || uncle.color == COLOR.BLANK
                    // 情况 3：
                    if(parent.left == cur){
                        rotateRight(parent);
                        RBTreeNode tmp = parent;
                        parent = cur;
                        cur = tmp;
                    }

                    // 情况 2：
                    rotateLeft(grandFather);
                    grandFather.color = COLOR.RED;
                    parent.color = COLOR.BLACK;
                }
            }
        }

        this.root.color = COLOR.BLACK;
        return true;
    }

    private void rotateLeft(RBTreeNode parent) {
        RBTreeNode subR = parent.right;
        RBTreeNode subRL = subR.left;

        parent.right = subRL;
        subR.left = parent;

        if(subRL != null){
            subRL.parent = parent;
        }

        RBTreeNode pParent = parent.parent;

        parent.parent = subR;

        if(root == parent){
            root = subR;
            root.parent = null;
        }else{
            if(pParent.left == parent){
                pParent.left = subR;
            }else{
                pParent.right = subR;
            }
        }

        subR.parent = pParent;
    }


    private void rotateRight(RBTreeNode parent) {
        RBTreeNode subL = parent.left;
        RBTreeNode subLR = subL.right;

        parent.left = subLR;
        subL.right = parent;

        if(subLR != null){
            subLR.parent = parent;
        }

        RBTreeNode pParent = parent.parent;
        parent.parent = subL;

        if(root == parent) {
            root = subL;
            root.parent = null;
        }else{
            if(pParent.left == parent){
                pParent.left = subL;
            }else{
                pParent.right = subL;
            }
        }


        subL.parent = pParent;
    }

    // 红黑树的验证分两步：

    // 第一步：验证是否为二叉搜索树
    public void inOrder(RBTreeNode root){
        if(root == null) return;

        inOrder(root.left);
        System.out.print(root.val + " ");
        inOrder(root.right);
    }

    // 第二步：验证是否符合红黑树的性质
    public boolean isRBTree(RBTreeNode root, int pathBlackNum, int blackNum){
        if(root == null) {
            return true;
        }
        // 验证红黑树的性质，分为三个性质进行验证
        // 第一步：验证根节点是否为黑色
        if(root.color == COLOR.RED){
            System.out.println("根节点为红色！");
            return false;
        }
        boolean flag = false;
        // 第二步：验证是否存在连续的两个红色的节点
        flag = checkRedNode(root);
        if(!flag) return false;
        // 第三步：验证是否每条路径上的黑色节点数量都相同
        flag = checkBlackNodeNum(root, pathBlackNum, blackNum);
        if(!flag) return false;

        return true;
    }

    private boolean checkRedNode(RBTreeNode root){
        if(root == null){
            return true;
        }
        // 遍历二叉树，如果遇到红色的节点，看一下它的父节点是否为红色即可
        if(root.color == COLOR.RED){
            RBTreeNode parent = root.parent;
            if(parent.color == COLOR.RED){
                System.out.println("连续两个红色的节点！");
                return false;
            }
        }
        return checkRedNode(root.left) && checkRedNode(root.right);
    }

    private boolean checkBlackNodeNum(RBTreeNode root, int pathBlackNum, int blackNum) {
        if(root == null){
            if(pathBlackNum == blackNum){
                return true;
            }
            System.out.println("路径上的黑色节点数量不相同！");
            return false;
        }

        if(root.color == COLOR.BLACK){
            pathBlackNum++;
        }

        return checkBlackNodeNum(root.left, pathBlackNum, blackNum) &&
                checkBlackNodeNum(root.right, pathBlackNum, blackNum);
    }
}
